package com.chen.system.controller;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.aigc.generation.models.QwenParam;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.MessageManager;
import com.alibaba.dashscope.common.ResultCallback;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.ApiException;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.alibaba.dashscope.utils.Constants;
import com.alibaba.dashscope.utils.JsonUtils;
import io.reactivex.Flowable;
import org.junit.jupiter.api.Test;
import org.springframework.boot.test.context.SpringBootTest;

import java.util.Arrays;
import java.util.Collections;
import java.util.concurrent.Semaphore;

//@SpringBootTest
class TongYiQianWenApplicationTests {

    static {
        Constants.apiKey = "sk-xxx";
    }


    @Test
    void contextLoads() {
    }

    /**
     * 通过messages调用
     *
     * @date 2024/02/23
     * @author Y.C
     */
    @Test
    void test001() {
        try {
            callWithMessage();
        } catch (ApiException | NoApiKeyException | InputRequiredException e) {
            System.out.println(e.getMessage());
        }
        System.exit(0);
    }

    /**
     * 通过prompt调用
     *
     * @date 2024/02/23
     * @author Y.C
     */
    @Test
    void test002() {
        try {
            qwenQuickStart();
            qwenQuickStartCallback();
        } catch (ApiException | NoApiKeyException | InputRequiredException | InterruptedException e) {
            System.out.printf("Exception %s%n", e.getMessage());
        }
        System.exit(0);
    }

    /**
     * 多轮会话 - 通过messages调用
     *
     * @date 2024/02/23
     * @author Y.C
     */
    @Test
    void test003() {
        try {
            callWithMessage02();
        } catch (ApiException | NoApiKeyException | InputRequiredException e) {
            System.out.println(e.getMessage());
        }
        System.exit(0);
    }

    /**
     * 流式输出
     *
     * @date 2024/02/23
     * @author Y.C
     */
    @Test
    void test004() {
        try {
            streamCallWithMessage();
        } catch (ApiException | NoApiKeyException | InputRequiredException e) {
            System.out.println(e.getMessage());
        }
        try {
            streamCallWithCallback();
        } catch (ApiException | NoApiKeyException | InputRequiredException | InterruptedException e) {
            System.out.println(e.getMessage());
        }
        System.exit(0);
    }





    /**
     * 通过messages调用
     *
     * @date 2024/02/23
     * @author Y.C
     */
    public static void callWithMessage()
            throws NoApiKeyException, ApiException, InputRequiredException {


        Generation gen = new Generation();
        MessageManager msgManager = new MessageManager(10);
        Message systemMsg =
                Message.builder().role(Role.SYSTEM.getValue()).content("You are a helpful assistant.").build();
        Message userMsg = Message.builder().role(Role.USER.getValue()).content("你好，周末去哪里玩？").build();
        msgManager.add(systemMsg);
        msgManager.add(userMsg);
        QwenParam param =
                QwenParam.builder().model(Generation.Models.QWEN_MAX).messages(msgManager.get())
                        .resultFormat(QwenParam.ResultFormat.MESSAGE)
                        .topP(0.8)
                        .enableSearch(true)
                        .build();
        GenerationResult result = gen.call(param);
        System.out.println(result);
        msgManager.add(result);
        System.out.println(JsonUtils.toJson(result));
        param.setPrompt("找个近点的");
        param.setMessages(msgManager.get());
        result = gen.call(param);
        System.out.println(result);
        System.out.println(JsonUtils.toJson(result));
    }


    /**
     * 提示
     */
    private final static String PROMPT = "如何做西红柿炖牛腩？";

    /**
     * 通过prompt调用
     *
     * @date 2024/02/23
     * @author Y.C
     */
    public static void qwenQuickStart()
            throws NoApiKeyException, ApiException, InputRequiredException {
        Generation gen = new Generation();
        QwenParam param = QwenParam.builder().model(Generation.Models.QWEN_TURBO).prompt(PROMPT)
                .topP(0.8).build();
        GenerationResult result = gen.call(param);
        System.out.println(JsonUtils.toJson(result));
    }

    /**
     * 通过prompt调用
     *
     * @date 2024/02/23
     * @author Y.C
     */
    public static void qwenQuickStartCallback()
            throws NoApiKeyException, ApiException, InputRequiredException, InterruptedException {
        Generation gen = new Generation();
        QwenParam param = QwenParam.builder().model(Generation.Models.QWEN_TURBO).prompt(PROMPT)
                .topP(0.8).build();
        Semaphore semaphore = new Semaphore(0);
        gen.call(param, new ResultCallback<GenerationResult>() {

            @Override
            public void onEvent(GenerationResult message) {
                System.out.println(message);
            }
            @Override
            public void onError(Exception ex){
                System.out.println(ex.getMessage());
                semaphore.release();
            }
            @Override
            public void onComplete(){
                System.out.println("onComplete");
                semaphore.release();
            }

        });
        semaphore.acquire();
    }


    /**
     * 多轮会话 - 通过messages调用
     *
     * @date 2024/02/23
     * @author Y.C
     */
    public static void callWithMessage02()
            throws NoApiKeyException, ApiException, InputRequiredException {
        Generation gen = new Generation();
        MessageManager msgManager = new MessageManager(10);
        Message systemMsg =
                Message.builder().role(Role.SYSTEM.getValue()).content("你是智能助手机器人").build();
        Message userMsg = Message.builder().role(Role.USER.getValue()).content("如何做西红柿炖牛腩？").build();
        msgManager.add(systemMsg);
        msgManager.add(userMsg);
        QwenParam param =
                QwenParam.builder().model(Generation.Models.QWEN_MAX).messages(msgManager.get())
                        .resultFormat(QwenParam.ResultFormat.MESSAGE)
                        .topP(0.8)
                        .enableSearch(true)
                        .build();
        GenerationResult result = gen.call(param);
        System.out.println(result);
        msgManager.add(result);
        System.out.println(JsonUtils.toJson(result));
        param.setPrompt("不放糖可以吗？");
        param.setMessages(msgManager.get());
        result = gen.call(param);
        System.out.println(result);
        System.out.println(JsonUtils.toJson(result));
    }


    /**
     * 使用消息流呼叫 - 流式输出
     *
     * @date 2024/02/23
     * @author Y.C
     */
    public static void streamCallWithMessage()
            throws NoApiKeyException, ApiException, InputRequiredException {
        Generation gen = new Generation();
        Message userMsg = Message
                .builder()
                .role(Role.USER.getValue())
                .content("如何做西红柿炖牛腩？")
                .build();
        QwenParam param =
                QwenParam.builder().model(Generation.Models.QWEN_MAX).messages(Collections.singletonList(userMsg))
                        .resultFormat(QwenParam.ResultFormat.MESSAGE)
                        .topP(0.8)
                        .enableSearch(true)
                        .incrementalOutput(true) // get streaming output incrementally
                        .build();
        Flowable<GenerationResult> result = gen.streamCall(param);
        StringBuilder fullContent = new StringBuilder();
        result.blockingForEach(message -> {
            fullContent.append(message.getOutput().getChoices().get(0).getMessage().getContent());
            System.out.println(JsonUtils.toJson(message));
        });
        System.out.println("Full content: \n" + fullContent.toString());
    }

    /**
     * 使用回调进行流式调用 - 流式输出
     *
     * @date 2024/02/23
     * @author Y.C
     */
    public static void streamCallWithCallback()
            throws NoApiKeyException, ApiException, InputRequiredException,InterruptedException {
        Generation gen = new Generation();
        Message userMsg = Message
                .builder()
                .role(Role.USER.getValue())
                .content("如何做西红柿炖牛腩？")
                .build();
        QwenParam param = QwenParam
                .builder()
                .model(Generation.Models.QWEN_MAX)
                .resultFormat(QwenParam.ResultFormat.MESSAGE)
                .messages(Collections.singletonList(userMsg))
                .topP(0.8)
                .incrementalOutput(true) // get streaming output incrementally
                .build();
        Semaphore semaphore = new Semaphore(0);
        StringBuilder fullContent = new StringBuilder();
        gen.streamCall(param, new ResultCallback<GenerationResult>() {

            @Override
            public void onEvent(GenerationResult message) {
                fullContent.append(message.getOutput().getChoices().get(0).getMessage().getContent());
                System.out.println(message);
            }
            @Override
            public void onError(Exception err){
                System.out.println(String.format("Exception: %s", err.getMessage()));
                semaphore.release();
            }

            @Override
            public void onComplete(){
                System.out.println("Completed");
                semaphore.release();
            }

        });
        semaphore.acquire();
        System.out.println("Full content: \n" + fullContent.toString());
    }
}